package com.shortlink.common.interceptor;

import com.alibaba.fastjson.JSON;
import com.google.common.collect.Lists;
import com.shortlink.common.constants.RedisConstant;
import com.shortlink.common.properties.UserFlowRiskControlProperties;
import com.shortlink.pojos.dto.dataobject.ResponseResult;
import com.shortlink.utils.UserContext;
import lombok.RequiredArgsConstructor;
import lombok.SneakyThrows;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.core.io.ClassPathResource;
import org.springframework.data.redis.core.StringRedisTemplate;
import org.springframework.data.redis.core.script.DefaultRedisScript;
import org.springframework.scripting.support.ResourceScriptSource;
import org.springframework.web.servlet.HandlerInterceptor;

import javax.servlet.*;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import java.io.IOException;
import java.io.PrintWriter;
import java.util.Optional;
import java.util.concurrent.TimeUnit;

import static com.shortlink.common.constants.AppHttpCodeEnum.FLOW_LIMIT_ERROR;

/**
 * 固定窗口限流
 */
@Slf4j
@RequiredArgsConstructor
public class UserFlowRiskControlInterceptor implements HandlerInterceptor {

    private final StringRedisTemplate stringRedisTemplate;
    private final UserFlowRiskControlProperties userFlowRiskControlProperties;

    private static final String USER_FLOW_RISK_CONTROL_LUA_SCRIPT_PATH = "lua/user_flow_risk_control.lua";

    @SneakyThrows
    @Override
    public boolean preHandle(HttpServletRequest request, HttpServletResponse response, Object handler) throws IOException, ServletException {
        DefaultRedisScript<Long> redisScript = new DefaultRedisScript<>();
        redisScript.setScriptSource(new ResourceScriptSource(new ClassPathResource(USER_FLOW_RISK_CONTROL_LUA_SCRIPT_PATH)));
        redisScript.setResultType(Long.class);
        String username = Optional.ofNullable(UserContext.getUsername()).orElse("other");
        Long result;
        try {
            // lua中timeWindow被双引号包围，由RedisConfig中value配置序列化器导致
            String timeWindow = userFlowRiskControlProperties.getTimeWindow();
            result = stringRedisTemplate.execute(redisScript, Lists.newArrayList(username), timeWindow);
        } catch (Throwable ex) {
            log.error("执行用户请求流量限制LUA脚本出错", ex);
            returnJson((HttpServletResponse) response, JSON.toJSONString(ResponseResult.errorResult(FLOW_LIMIT_ERROR)));
            return false;
        }
        if (result == null || result > userFlowRiskControlProperties.getMaxAccessCount()) {
            returnJson((HttpServletResponse) response, JSON.toJSONString(ResponseResult.errorResult(FLOW_LIMIT_ERROR)));

            // 拉黑
            String clientIP = request.getRemoteAddr();
            String key= RedisConstant.BLACKLIST_KEY_PREFIX +clientIP;
            stringRedisTemplate.opsForValue().set(key, "1", 5, TimeUnit.SECONDS);

            return false;
        }
        return true;
    }

    private void returnJson(HttpServletResponse response, String json) throws Exception {
        response.setCharacterEncoding("UTF-8");
        response.setContentType("text/html; charset=utf-8");
        try (PrintWriter writer = response.getWriter()) {
            writer.print(json);
        }
    }
}

